import pandas as pd
import time
import threading
from consts import (
  FEATURE_KEYWORD,
  FEATURE_GENRE,
  FEATURE_INFO,
  FEATURE_ALL,

  PLATFORMS_BAIDU,
  PLATFORMS_ALI,

  # 百度大模型
  QIANFAN_ERNIE_SPEED_8K,
  QIANFAN_ERNIE_SPEED_128K,
  QIANFAN_ERNIE_LITE_8K,
  QIANFAN_ERNIE_LITE_8K_0922,
  QIANFAN_ERNIE_LITE_8K_0308,
  QIFAFAN_ERNIE_TINY_8K,

  # 阿里大模型
  BAILIAN_QWEN_MAX,
  BAILIAN_QWEN_MAX_0428,
  BAILIAN_QWEN_LONG,
  BAILIAN_QWEN2_72B_INSTRUCT,
  BAILIAN_QWEN2_57B_A14B_INSTRUCT,
  BAILIAN_QWEN2_7B_INSTRUCT,
  BAILIAN_QWEN2_1_5B_INSTRUCT,
  BAILIAN_QWEN2_0_5B_INSTRUCT,
  BAILIAN_QWEN_MAX_LONG_CONTEXT,
  BAILIAN_QWEN_PLUS,
  BAILIAN_QWEN_TURBO,
  BAILIAN_QWEN_1_5_110B_CHAT,
)
from utils import (
  log_message
)

from get_data_and_prompt import (
  get_train_test_data,
  create_prompt
)

from qianfan_model import (
  run_qianfan_model
)
from qwen_model import (
  run_bailian_model
)

# 构造提示
def run_model(feature_type, platform, model_name, random_state=42):
  log_message(f"正在使用{feature_type}特征值相关数据运行{model_name}模型")
  data = get_train_test_data(feature_type)
  train_data = data["train_data"]
  test_data = data["test_data"]
  output = []
  count = 0
  for user_id in train_data["userId"].unique():
    count += 1
    user_train_data = train_data[train_data["userId"] == user_id]
    user_test_data = test_data[test_data["userId"] == user_id]
    train_records = user_train_data.to_dict(orient='records')
    test_records = user_test_data.to_dict(orient='records')

    test_original_movie_id = test_records[0]["originalMovieId"]

    prompt = create_prompt(feature_type, train_records, test_records)
    system_prompt = prompt["system_prompt"]
    full_prompt = prompt["full_prompt"]
    circle_time = 0
    # 重复捕获异常
    while True:
      circle_time += 1
      if circle_time > 10:
        log_message(f"发生错误超过10次，跳过该用户,默认取2.5评分")
        row = {"userId": user_id, "movieId": test_original_movie_id, "predict": "2.5"}
        output.append(row)
        # log_message(row)
        break
      try:
        # 用模型，参数为：system_prompt, full_prompt
        if platform == PLATFORMS_BAIDU:
          result = run_qianfan_model(system_prompt, full_prompt, model_name)
        elif platform == PLATFORMS_ALI:
          result = run_bailian_model(system_prompt, full_prompt, model_name)

        row = {"userId": user_id, "movieId": test_original_movie_id, "predict": str(result)}
        output.append(row)
        # log_message(f"{model_name}模型，数据{row}")
        # 如果成功跳出循环
        break
      except Exception as e:
        log_message(f"发生值错误{e}等待20秒后再继续请求")
        time.sleep(20)

    # 输出运行进度
    if count % 50 == 0:
      log_message(f"{model_name}模型运行进度：{round(count/len(train_data['userId'].unique()), 2) * 100} %")
  csv_file_path = f"../data/result/basic_large_model/origin_data/{feature_type}/{model_name}_random_state_{random_state}_1k.csv"
  log_message(f"数据{csv_file_path}存储完毕")
  result = pd.DataFrame(output)
  result.to_csv(csv_file_path, index=False)



# 创建并启动线程的函数，接受目标函数和参数
def start_thread(target_function, **kwargs):
  thread = threading.Thread(target=target_function, kwargs=kwargs)
  thread.start()
  return thread

# 百度千帆模型多线程跑
def run_qianfan_model_feature_all_thread():
  # 多线程启动
  threads = [
    # all
    start_thread(run_model, feature_type=FEATURE_ALL, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_SPEED_8K),
    start_thread(run_model, feature_type=FEATURE_ALL, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_SPEED_128K),
    start_thread(run_model, feature_type=FEATURE_ALL, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K),
    start_thread(run_model, feature_type=FEATURE_ALL, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K_0922),
    start_thread(run_model, feature_type=FEATURE_ALL, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K_0308),
    start_thread(run_model, feature_type=FEATURE_ALL, platform=PLATFORMS_BAIDU, model_name=QIFAFAN_ERNIE_TINY_8K),
    # info
    start_thread(run_model, feature_type=FEATURE_INFO, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_SPEED_8K),
    start_thread(run_model, feature_type=FEATURE_INFO, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_SPEED_128K),
    start_thread(run_model, feature_type=FEATURE_INFO, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K),
    start_thread(run_model, feature_type=FEATURE_INFO, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K_0922),
    start_thread(run_model, feature_type=FEATURE_INFO, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K_0308),
    start_thread(run_model, feature_type=FEATURE_INFO, platform=PLATFORMS_BAIDU, model_name=QIFAFAN_ERNIE_TINY_8K),
    # genre
    start_thread(run_model, feature_type=FEATURE_GENRE, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_SPEED_8K),
    start_thread(run_model, feature_type=FEATURE_GENRE, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_SPEED_128K),
    start_thread(run_model, feature_type=FEATURE_GENRE, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K),
    start_thread(run_model, feature_type=FEATURE_GENRE, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K_0922),
    start_thread(run_model, feature_type=FEATURE_GENRE, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K_0308),
    start_thread(run_model, feature_type=FEATURE_GENRE, platform=PLATFORMS_BAIDU, model_name=QIFAFAN_ERNIE_TINY_8K),
    # keyword
    start_thread(run_model, feature_type=FEATURE_KEYWORD, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_SPEED_8K),
    start_thread(run_model, feature_type=FEATURE_KEYWORD, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_SPEED_128K),
    start_thread(run_model, feature_type=FEATURE_KEYWORD, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K),
    start_thread(run_model, feature_type=FEATURE_KEYWORD, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K_0922),
    start_thread(run_model, feature_type=FEATURE_KEYWORD, platform=PLATFORMS_BAIDU, model_name=QIANFAN_ERNIE_LITE_8K_0308),
    start_thread(run_model, feature_type=FEATURE_KEYWORD, platform=PLATFORMS_BAIDU, model_name=QIFAFAN_ERNIE_TINY_8K),
  ]
  for thread in threads:
    thread.join()
  log_message("所有线程程序完成")

run_qianfan_model_feature_all_thread()
